#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SM4特定的Side-Channel Attack
"""

import numpy as np
import sys
from pathlib import Path
from .sm4_models import SM4LeakageModels
from .sm4_key_schedule_reverse import SM4KeyScheduleReverse
from ..base_algorithm_attack import AlgorithmSpecificAttack
from typing import List, Dict, Any
import logging

# Createlogger（NotUsebasicConfig，Avoid重复日志）
logger = logging.getLogger('AlgorithmAttack.SM4')


class SM4SpecificAttacks(AlgorithmSpecificAttack):
    """SM4特定的Side-Channel Attack"""
    
    def __init__(self):
        super().__init__("SM4")
        
        # SM4Parameters
        self.key_size = 16  # SM4Key长度
        self.block_size = 16  # SM4分组长度
        self.round_count = 32  # SM4轮数
        
        # 新架构：分离的模型
        self.intermediate_functions = SM4LeakageModels.get_intermediate_functions()
        self.leakage_models = SM4LeakageModels.get_leakage_models()
        self.operation_mapping = SM4LeakageModels.get_operation_to_intermediate_mapping()
        
        # Save目标字节索引
        self._target_byte_cache = 0
        
    def get_algorithm_parameters(self) -> Dict[str, Any]:
        """获取SM4AlgorithmParameters"""
        return {
            'algorithm': 'SM4',
            'type': 'symmetric',
            'key_size': self.key_size,
            'block_size': self.block_size,
            'round_count': self.round_count,
            'rounds': self.round_count,
            'sbox_size': 256,
            'word_size': 32
        }
    
    def get_intermediate_functions(self) -> Dict[str, Any]:
        """获取Intermediate valueCalculateFunction"""
        return self.intermediate_functions
    
    def get_leakage_models(self) -> Dict[str, Any]:
        """获取泄漏表示模型"""
        return self.leakage_models
    
    def create_leakage_models(self) -> Dict[str, Any]:
        """CreateSM4特定的泄漏模型（用于基Class兼容）"""
        combined = {}
        combined.update(self.intermediate_functions)
        combined.update(self.leakage_models)
        return combined
    
    def get_target_operations(self) -> List[str]:
        """获取SM4可Attack的目标操作"""
        return SM4LeakageModels.get_target_operations()
    
    def get_operation_to_intermediate_mapping(self) -> Dict[str, str]:
        """获取目标操作ToIntermediate valueFunction的Mapping"""
        return self.operation_mapping
    
    def _calculate_intermediate_values(self, target_data: np.ndarray, hypothesis: int, 
                                     target_operation: str) -> np.ndarray:
        """
        CalculateSM4特定的Intermediate value
        
        Args:
            target_data: 目标Data（通常YesPlaintext的某个字节）
            hypothesis: False设Key字节
            target_operation: 目标操作
            
        Returns:
            Intermediate valueArray（汉明权重）
        """
        if target_operation == 'sbox_output' or target_operation == 'first_round' or target_operation == 'sbox':
            # 标准CPAAttack模型：UseS-boxOutput的汉明权重
            # 理论基础：Power与Data的汉明权重成正比
            
            results = []
            for td in target_data:
                pt_byte = int(td) & 0xFF  # Plaintext字节
                key_byte = hypothesis & 0xFF  # Key字节False设
                
                # S-boxInput: plaintext XOR key
                sbox_input = pt_byte ^ key_byte
                
                # S-boxOutput
                sbox_output = SM4LeakageModels.sbox_output_model(sbox_input)
                
                # 标准泄漏模型：S-boxOutput的汉明权重
                # 这YesCPAAttack的标准做法
                hw_output = bin(sbox_output).count('1')
                results.append(hw_output)
            
            return np.array(results)
        
        elif target_operation == 'round_function':
            # 轮Function泄漏模型
            sbox_outputs = np.array([SM4LeakageModels.sbox_output_model(pt ^ hypothesis) 
                           for pt in target_data])
            return np.array([SM4LeakageModels.hamming_weight(out) for out in sbox_outputs])
        
        elif target_operation == 'linear_transform':
            # 线性变换泄漏模型（Complete32位实现）
            # L(B) = B XOR (B<<<2) XOR (B<<<10) XOR (B<<<18) XOR (B<<<24)
            results = []
            for td in target_data:
                pt_byte = int(td) & 0xFF
                key_byte = hypothesis & 0xFF
                
                # S-boxInputAndOutput
                sbox_input = pt_byte ^ key_byte
                sbox_output = SM4LeakageModels.sbox_output_model(sbox_input)
                
                # 扩展To32位（False设Yes4个相同字节，实际应According to上下文）
                sbox_32bit = (sbox_output << 24) | (sbox_output << 16) | (sbox_output << 8) | sbox_output
                
                # 应用Complete的L变换
                transformed = self._linear_transform_l(sbox_32bit)
                
                # Return汉明权重
                hw = bin(transformed).count('1')
                results.append(hw)
            
            return np.array(results)
        
        elif target_operation == 'round_key_xor':
            # 轮KeyXOR
            xor_values = target_data ^ hypothesis
            return np.array([SM4LeakageModels.hamming_weight(val) for val in xor_values])
        
        else:
            # DefaultUseS-boxOutput模型
            sbox_outputs = np.array([SM4LeakageModels.sbox_output_model(pt ^ hypothesis) 
                           for pt in target_data])
            return np.array([SM4LeakageModels.hamming_weight(out) for out in sbox_outputs])
    
    def _calculate_sm4_round_hw(self, plaintext: bytes, hypothesis: int, round_num: int = 0,
                                 recovered_previous_rks: List[int] = None) -> int:
        """
        CalculateSM4指定轮次的汉明权重泄漏（通用多轮Support）
        
        Args:
            plaintext: 16字节Plaintext
            hypothesis: Current轮Key的某个字节False设 (0-255)
            round_num: 轮次索引 (0-31)
            recovered_previous_rks: 已恢复的前面轮次的轮KeyList (32位Integer)
            
        Returns:
            S-boxOutput的汉明权重
        """
        # 将PlaintextConvertFor4个32位字
        X = [int.from_bytes(plaintext[i*4:(i+1)*4], 'big') for i in range(4)]
        
        # Execute前面的轮次（Use已恢复的轮Key）
        for r in range(round_num):
            if recovered_previous_rks is None or r >= len(recovered_previous_rks):
                # If没HasProvide前面的轮Key，No法Calculate
                raise ValueError(f"Require前{round_num}个轮Key才能Attack第{round_num}轮")
            
            rk = recovered_previous_rks[r]
            temp = X[1] ^ X[2] ^ X[3] ^ rk
            
            # Complete的SM4轮Function T = L(τ(temp))
            # 1. 非线性变换τ（4个并行S-box）
            t_tau = 0
            for i in range(4):
                byte_val = (temp >> (24 - i * 8)) & 0xFF
                sbox_out = SM4LeakageModels.sbox_output_model(byte_val)
                t_tau |= (sbox_out << (24 - i * 8))
            
            # 2. 线性变换L [OK] 这Yes关键修复！
            t_result = self._linear_transform_l(t_tau)
            
            # 3. Calculate新Status
            X_new = X[0] ^ t_result
            X[0], X[1], X[2], X[3] = X[1], X[2], X[3], X_new
        
        # CalculateCurrent轮的temp值（NotInclude轮Key）
        temp_without_key = X[1] ^ X[2] ^ X[3]
        
        # 确定目标字节位置
        temp_byte_idx = self._target_byte_cache % 4
        
        # Extract目标字节
        temp_byte_no_key = (temp_without_key >> (24 - temp_byte_idx * 8)) & 0xFF
        
        # 加上False设的轮Key字节
        sbox_input = temp_byte_no_key ^ hypothesis
        
        # S-box变换
        sbox_output = SM4LeakageModels.sbox_output_model(sbox_input)
        
        # Return汉明权重
        return SM4LeakageModels.hamming_weight(sbox_output)
    
    def _linear_transform_l(self, x: int) -> int:
        """
        SM4线性变换L
        L(B) = B ⊕ (B<<<2) ⊕ (B<<<10) ⊕ (B<<<18) ⊕ (B<<<24)
        
        Args:
            x: 32位Integer
            
        Returns:
            线性变换后的32位Integer
        """
        def rotl(value: int, shift: int) -> int:
            """32位Loop左移"""
            # EnsureUsePython int，AvoidnumpyClass型问题
            value = int(value) & 0xFFFFFFFF
            return int(((value << shift) | (value >> (32 - shift))) & 0xFFFFFFFF)
        
        x = int(x) & 0xFFFFFFFF  # EnsureYesPython int
        result = x
        result ^= rotl(x, 2)
        result ^= rotl(x, 10)
        result ^= rotl(x, 18)
        result ^= rotl(x, 24)
        return int(result & 0xFFFFFFFF)
    
    
    def _get_selection_function(self, target_data: np.ndarray, hypothesis: int,
                              target_operation: str) -> np.ndarray:
        """
        获取SM4特定的选择Function（用于DPAAttack）
        
        Args:
            target_data: 目标Data
            hypothesis: False设Key字节
            target_operation: 目标操作
            
        Returns:
            选择FunctionResult（0Or1的Array）
        """
        intermediate_values = self._calculate_intermediate_values(
            target_data, hypothesis, target_operation
        )
        
        if target_operation in ['sbox_output', 'first_round']:
            # UseS-boxOutput的最低位作For选择Function
            return intermediate_values & 1
        
        elif target_operation == 'round_function':
            # Use轮FunctionOutput的汉明权重奇偶性
            hamming_weights = np.array([SM4LeakageModels.hamming_weight(val) 
                                      for val in intermediate_values])
            return hamming_weights & 1
        
        elif target_operation == 'linear_transform':
            # Use线性变换Output的最高位
            return (intermediate_values >> 7) & 1
        
        else:
            # DefaultUse最低位
            return intermediate_values & 1
    
    def attack_round_sbox(self, round_num: int = 0, target_byte: int = 0, 
                         recovered_previous_rks: List[int] = None, **kwargs):
        """
        Attack指定轮次的S-box（通用多轮Support）
        
        Args:
            round_num: 轮次索引 (0-31)
            target_byte: 目标字节位置 (0-3，对应rk[round_num]的字节索引)
            recovered_previous_rks: 已恢复的前面轮次的Complete轮KeyList (32位Integer)
            **kwargs: AttackParameters
            
        Returns:
            AttackResult（recovered_key_bytes中YesCurrent轮的轮Key字节）
        """
        # Save轮次InformationAnd已恢复的轮Key
        self._current_round_num = round_num
        self._recovered_previous_rks = recovered_previous_rks if recovered_previous_rks else []
        
        # Save目标字节索引（自动取模To0-3范围）
        self._target_byte_cache = target_byte % 4
        
        return self.run_cpa_attack(
            target_operation='first_round',  # Usefirst_round操作
            target_byte=target_byte,
            **kwargs
        )
    
    def attack_first_round_sbox(self, target_byte: int = 0, **kwargs):
        """
        AttackFirst轮S-box（最常见的SM4Attack）
        
        WARNING - Important：
        此Attack恢复的Yes轮Keyrk[0]的字节，而NotYes主Key的字节！
        
        SM4First轮：X[4] = X[0] ⊕ T(X[1] ⊕ X[2] ⊕ X[3] ⊕ rk[0])
        泄漏点：temp = X[1] ⊕ X[2] ⊕ X[3] ⊕ rk[0]
        
        由于rk[0]Yes32位（4字节），target_byte的Has效范围Yes0-3，对应rk[0]的4个字节。
        Iftarget_byte >= 4，会自动取模：target_byte % 4
        
        Args:
            target_byte: 目标字节位置 (0-3，对应rk[0]的字节索引)
            **kwargs: AttackParameters
            
        Returns:
            AttackResult（recovered_key_bytes中Yes轮Key字节，NotYes主Key字节）
        """
        # Save目标字节索引供泄漏模型Use（自动取模To0-3范围）
        self._target_byte_cache = target_byte % 4
        
        return self.run_cpa_attack(
            target_operation='first_round',
            target_byte=target_byte,
            **kwargs
        )
    
    def attack_key_schedule(self, **kwargs):
        """
        AttackKey调度过程
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_spa_attack(
            target_operation='key_schedule',
            **kwargs
        )
    
    def attack_linear_transform(self, **kwargs):
        """
        Attack线性变换L操作
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_template_attack(
            target_operation='linear_transform',
            **kwargs
        )
    
    def full_key_recovery_attack(self, **kwargs) -> Dict[str, Any]:
        """
        SM4First轮轮Key恢复Attack
        
        WARNING - Important：
        此Attack恢复的Yes轮Keyrk[0]（32位，4字节），而NotYes主Key（128位，16字节）！
        
        SM4的Key调度Algorithm非常复杂，CPAAttackFirst轮只能恢复rk[0]的4个字节。
        IfRequire恢复主Key，Require：
        1. AttackMultiple轮次（至少前8轮）
        2. 实现SM4Key调度的逆向Algorithm
        3. FromMultiple轮Key推导主Key
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            CompleteAttackResult（Includerk[0]的4个字节）
        """
        results = {}
        recovered_rk0_bytes = {}
        
        self.logger.info("Starting SM4 round key rk[0] recovery attack")
        self.logger.info("WARNING: Attacking round key rk[0] (4 bytes), not master key (16 bytes)!")
        
        # 只Attackrk[0]的4个字节
        for byte_pos in range(4):
            self.logger.info(f"Attacking rk[0] byte {byte_pos}")
            
            try:
                # UseCPAAttackFirst轮S-box
                result = self.attack_first_round_sbox(
                    target_byte=byte_pos,
                    **kwargs
                )
                
                results[f"rk0_byte_{byte_pos}"] = result
                
                if result.success and result.recovered_key_bytes:
                    recovered_rk0_bytes[byte_pos] = result.recovered_key_bytes[0]
                    self.logger.info(f"rk[0] Byte {byte_pos}: 0x{result.recovered_key_bytes[0]:02X} "
                                   f"(confidence: {result.confidence:.4f})")
                else:
                    self.logger.warning(f"Failed to recover rk[0] byte {byte_pos}")
            
            except Exception as e:
                self.logger.error(f"Error attacking rk[0] byte {byte_pos}: {e}")
        
        # 构建恢复的rk[0]
        recovered_rk0 = []
        for i in range(4):
            if i in recovered_rk0_bytes:
                recovered_rk0.append(recovered_rk0_bytes[i])
            else:
                recovered_rk0.append(0x00)  # 未恢复的字节设For0
        
        # 将rk[0]的4个字节组合成32位Integer
        recovered_rk0_value = (recovered_rk0[0] << 24) | (recovered_rk0[1] << 16) | \
                              (recovered_rk0[2] << 8) | recovered_rk0[3]
        
        success_rate = len(recovered_rk0_bytes) / 4
        
        summary = {
            'algorithm': 'SM4',
            'attack_type': 'Round Key rk[0] Recovery',
            'target': 'Round Key rk[0] (not master key!)',
            'total_bytes': 4,  # rk[0]只Has4个字节
            'recovered_bytes': len(recovered_rk0_bytes),
            'success_rate': success_rate,
            'recovered_round_key_bytes': recovered_rk0,
            'recovered_round_key_hex': ''.join(f'{b:02X}' for b in recovered_rk0),
            'recovered_round_key_value': f"0x{recovered_rk0_value:08X}",
            'individual_results': results,
            'note': 'This is the round key rk[0], not the master key. ' +
                   'To recover the master key, you need to attack multiple rounds ' +
                   'and implement SM4 key schedule inversion.'
        }
        
        self.logger.info(f"SM4 round key recovery completed: {len(recovered_rk0_bytes)}/4 bytes recovered "
                        f"({success_rate:.1%} success rate)")
        self.logger.info(f"Recovered rk[0]: {summary['recovered_round_key_value']}")
        
        return summary
    
    def full_master_key_recovery_attack(self, **kwargs) -> Dict[str, Any]:
        """
        Complete的SM4主Key恢复Attack（16字节）
        
        策略：
        1. Attack前4轮，恢复rk[0], rk[1], rk[2], rk[3]
        2. UseKey调度逆向Algorithm，From4个轮Key推导主Key
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            CompleteAttackResult（Include恢复的16字节主Key）
        """
        self.logger.info("="*80)
        self.logger.info("Starting SM4 FULL MASTER KEY recovery attack (16 bytes)")
        self.logger.info("="*80)
        self.logger.info("Strategy: Attack 4 rounds -> Recover rk[0-3] -> Reverse key schedule -> Master key")
        
        results = {}
        recovered_round_keys = []  # 存储Complete的32位轮Key
        recovered_round_key_bytes = {}  # 存储Each字节的恢复Result
        
        # 阶段1：Attack前4轮，每轮恢复4个字节
        for round_num in range(4):
            self.logger.info(f"\n{'='*80}")
            self.logger.info(f"Round {round_num}: Attacking rk[{round_num}]")
            self.logger.info(f"{'='*80}")
            
            round_key_bytes = {}
            
            for byte_pos in range(4):
                self.logger.info(f"  Attacking rk[{round_num}] byte {byte_pos}...")
                
                try:
                    # Use通用的多轮AttackMethod
                    result = self.attack_round_sbox(
                        round_num=round_num,
                        target_byte=byte_pos,
                        recovered_previous_rks=recovered_round_keys,
                        **kwargs
                    )
                    
                    # 存储可序列化的Result
                    results[f"rk{round_num}_byte_{byte_pos}"] = {
                        'recovered_byte': int(result.recovered_key_bytes[0]) if result.recovered_key_bytes else 0,
                        'confidence': float(result.confidence),
                        'success': bool(result.success)
                    }
                    
                    if result.success and result.recovered_key_bytes:
                        byte_value = result.recovered_key_bytes[0]
                        round_key_bytes[byte_pos] = byte_value
                        recovered_round_key_bytes[f"rk{round_num}_byte{byte_pos}"] = byte_value
                        
                        self.logger.info(f"  [OK] rk[{round_num}] Byte {byte_pos}: 0x{byte_value:02X} "
                                       f"(confidence: {result.confidence:.4f})")
                    else:
                        self.logger.warning(f"  [FAIL] Failed to recover rk[{round_num}] byte {byte_pos}")
                        round_key_bytes[byte_pos] = 0x00
                
                except Exception as e:
                    self.logger.error(f"  [ERROR] Error attacking rk[{round_num}] byte {byte_pos}: {e}")
                    round_key_bytes[byte_pos] = 0x00
                    results[f"rk{round_num}_byte_{byte_pos}"] = {
                        'recovered_byte': 0,
                        'confidence': 0.0,
                        'success': False
                    }
            
            # 将4个字节组合成32位轮Key
            if len(round_key_bytes) == 4:
                rk_value = (round_key_bytes[0] << 24) | (round_key_bytes[1] << 16) | \
                           (round_key_bytes[2] << 8) | round_key_bytes[3]
                recovered_round_keys.append(rk_value)
                self.logger.info(f"[OK] Round {round_num} completed: rk[{round_num}] = 0x{rk_value:08X}")
            else:
                self.logger.error(f"[FAIL] Round {round_num} incomplete, using 0x00000000")
                recovered_round_keys.append(0x00000000)
        
        # 阶段2：From轮Key逆推主Key
        self.logger.info(f"\n{'='*80}")
        self.logger.info("Phase 2: Reverse key schedule to recover master key")
        self.logger.info(f"{'='*80}")
        
        master_key = None
        master_key_hex = None
        reverse_success = False
        
        if len(recovered_round_keys) == 4:
            try:
                self.logger.info(f"Recovered round keys:")
                for i, rk in enumerate(recovered_round_keys):
                    self.logger.info(f"  rk[{i}] = 0x{rk:08X}")
                
                # UseKey调度逆向Algorithm
                master_key = SM4KeyScheduleReverse.reverse_from_first_4_round_keys(
                    recovered_round_keys[0],
                    recovered_round_keys[1],
                    recovered_round_keys[2],
                    recovered_round_keys[3]
                )
                
                master_key_hex = master_key.hex()
                
                # Verify恢复的主Key
                is_valid = SM4KeyScheduleReverse.verify_master_key(
                    master_key,
                    recovered_round_keys[0],
                    recovered_round_keys[1],
                    recovered_round_keys[2],
                    recovered_round_keys[3]
                )
                
                if is_valid:
                    self.logger.info(f"[OK] Master key recovered and verified!")
                    self.logger.info(f"   Master Key: {master_key_hex}")
                    reverse_success = True
                else:
                    self.logger.warning(f"[WARNING] Master key recovered but verification failed")
                    reverse_success = False
                    
            except Exception as e:
                self.logger.error(f"[FAIL] Failed to reverse key schedule: {e}")
                import traceback
                traceback.print_exc()
        else:
            self.logger.error(f"[FAIL] Cannot reverse key schedule: need 4 round keys, got {len(recovered_round_keys)}")
        
        # 构建Final result
        success_rate = len([v for v in recovered_round_key_bytes.values() if v != 0x00]) / 16
        
        summary = {
            'algorithm': 'SM4',
            'attack_type': 'Full Master Key Recovery (Multi-Round Attack)',
            'total_bytes': int(16),
            'recovered_bytes': int(len([v for v in recovered_round_key_bytes.values() if v != 0x00])),
            'success_rate': float(success_rate),
            'rounds_attacked': int(4),
            'recovered_round_keys': [str(f"0x{int(rk):08X}") for rk in recovered_round_keys],
            'recovered_round_key_bytes': {str(k): int(v) for k, v in recovered_round_key_bytes.items()},
            'master_key_recovered': bool(reverse_success),
            'master_key_hex': str(master_key_hex) if master_key_hex else 'N/A',
            'master_key_bytes': [int(b) for b in master_key] if master_key else None,
            'individual_results': results,  # 已经Yes可序列化的Format
            'note': 'Multi-round attack: recovered rk[0-3] and reversed key schedule to obtain master key'
        }
        
        self.logger.info(f"\n{'='*80}")
        self.logger.info("SM4 Master Key Recovery Attack Completed")
        self.logger.info(f"{'='*80}")
        self.logger.info(f"Success rate: {success_rate:.1%}")
        if reverse_success:
            self.logger.info(f"[OK] Master Key: {master_key_hex}")
        else:
            self.logger.info(f"[FAIL] Master key recovery failed")
        
        return summary
    
    def validate_recovered_key(self, recovered_key: List[int], 
                             test_plaintext: bytes, test_ciphertext: bytes) -> bool:
        """
        Verify恢复的KeyYesNo正确
        
        Args:
            recovered_key: 恢复的Key字节List
            test_plaintext: TestPlaintext
            test_ciphertext: TestCiphertext
            
        Returns:
            KeyYesNo正确
        """
        try:
            # CallSM4加密来Verify
            from ...crypto.sm4_cipher import SM4Cipher
            
            key_bytes = bytes(recovered_key)
            sm4 = SM4Cipher(key_bytes)
            
            # 加密TestPlaintext
            encrypted = sm4.encrypt(test_plaintext)
            
            # CompareResult
            return encrypted == test_ciphertext
            
        except Exception as e:
            self.logger.error(f"Key validation failed: {e}")
            return False
